{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import lasagne\n",
    "import theano\n",
    "from theano import tensor as T\n",
    "import sys\n",
    "import numpy as np\n",
    "import json\n",
    "import pickle\n",
    "import os\n",
    "#enable importing of notebooks\n",
    "from nbfinder import NotebookFinder\n",
    "sys.meta_path.append(NotebookFinder())\n",
    "# from print_n_plot import plot_ims_with_boxes, add_bbox, plot_im_with_box\n",
    "from copy import deepcopy\n",
    "from collections import Counter\n",
    "from os.path import join, exists\n",
    "from os import mkdir, makedirs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def makedir_if_not_there(dirname):\n",
    "    if not exists(dirname):\n",
    "        try:\n",
    "            mkdir(dirname)\n",
    "        except OSError:\n",
    "            makedirs(dirname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def softmax3D(x):\n",
    "    #X is batch_size x n_classes x X x Y tensor\n",
    "    s = x.shape\n",
    "    #Flip X to be batch_size*x*y, n_classes -> stack all the one-hot encoded vectors\n",
    "    x_t = x.transpose((0,2,3,1)).reshape((s[0]*s[2]*s[3],s[1]))\n",
    "    # take softmax        \n",
    "    x_sm = T.nnet.softmax(x_t)\n",
    "    #reshape back to #X is batch_size x n_classes x X x Y tensor\n",
    "    x_f = x_sm.reshape((s[0],s[2],s[3],s[1])).transpose((0,3,1,2))\n",
    "    return x_f\n",
    "\n",
    "def softmax4D(x):\n",
    "    #X is batch_size x n_classes x time-steps x X x Y tensor\n",
    "    s = x.shape\n",
    "    #Flip X to be batch_size*x*y, n_classes -> stack all the one-hot encoded vectors\n",
    "    x_t = x.transpose((0,2,3,4,1)).reshape((s[0]*s[2]*s[3]*s[4],s[1]))\n",
    "    # take softmax        \n",
    "    x_sm = T.nnet.softmax(x_t)\n",
    "    #reshape back to #X is batch_size x n_classes x X x Y tensor\n",
    "    x_f = x_sm.reshape((s[0],s[2],s[3],s[4],s[1])).transpose((0,4,1,2,3))\n",
    "    return x_f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def smoothL1(x):\n",
    "    #x is vector of scalars\n",
    "    lto = T.abs_(x)<1\n",
    "    gteo = T.abs_(x)>=1\n",
    "    new_x = T.set_subtensor(x[lto.nonzero()],0.5 * T.square(x[lto.nonzero()]))\n",
    "    new_x = T.set_subtensor(new_x[gteo.nonzero()], T.abs_(new_x[gteo.nonzero()]) - 0.5)\n",
    "    return new_x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def add_as_appension(key, val, dict_,):\n",
    "    if key not in dict_:\n",
    "        dict_[key] = []\n",
    "    dict_[key].append(val)\n",
    "    return dict_\n",
    "\n",
    "def add_as_extension(key, val, dict_):\n",
    "    if key not in dict_:\n",
    "        dict_[key] = []\n",
    "    dict_[key].extend(val)\n",
    "    return dict_\n",
    "\n",
    "def add_as_running_total(key,val,dict_):\n",
    "    if key not in dict_:\n",
    "        dict_[key] = 0\n",
    "    dict_[key] += val\n",
    "    return dict_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class AccuracyGetter(object):\n",
    "    def __init__(self,kwargs):\n",
    "        self.kwargs = kwargs\n",
    "        self.class_name_dict = dict(zip(range(4),[\"td\", \"tc\", \"etc\", \"ar\"], ))\n",
    "\n",
    "    def get_scores(self, pred_tensor,y_tensor):\n",
    "        pred_confs_tot = {}\n",
    "        gt_confs_tot = {\"gt_\" + str(k):[] for k in range(self.kwargs[\"num_classes\"])}\n",
    "        pred_confs_tot = {\"pred_\" + str(k) :[] for k in range(self.kwargs[\"num_classes\"])}\n",
    "        for i in range(pred_tensor.shape[0]):\n",
    "\n",
    "            pred_boxes, gt_boxes = self.get_boxes(pred_tensor[i], y_tensor[i])\n",
    "            pred_confs, gt_confs = self.get_pred_and_gt_scores(pred_boxes,gt_boxes)\n",
    "            for k in pred_confs.keys():\n",
    "                pred_confs_tot[\"pred_\" + str(k)].extend(pred_confs[k])\n",
    "                gt_confs_tot[\"gt_\" + str(k)].extend(gt_confs[k])\n",
    "                \n",
    "        pred_confs_tot.update(gt_confs_tot)\n",
    "        return pred_confs_tot\n",
    "            \n",
    "\n",
    "\n",
    "#     def get_ap(self, pred_boxes,gt_boxes):\n",
    "#         '''pred boxes and gt_boxes is a dictionary key (class integer): value: list of boxes'''\n",
    "        \n",
    "#         aps = {k: average_precision_score(gt_confs[k], pred_confs[k])}\n",
    "    \n",
    " \n",
    "    def get_pred_and_gt_scores(self, pred_boxes,gt_boxes):\n",
    "        '''input: bbox lists for predicted and ground truth (gt) boxes\n",
    "           each element in list is a list -> [x, y, w, h, confidence, class]\n",
    "        output: \n",
    "                pred_confs: dictionary, where keys are classes and values are list confidences for each bounding box for that class\n",
    "                gt_confs: dictionary where keys are classes and values are list of binary variables, corresponding to crrectness of \n",
    "                    each element in corresponding lists in pred_confs'''\n",
    "                \n",
    "        \n",
    "        \n",
    "        #dict where key is class, value is list of boxes predicted as that class\n",
    "        pred_boxes_by_class = {k: [p for p in pred_boxes if p[5] ==k] for k in range(self.kwargs[\"num_classes\"])}\n",
    "        \n",
    "        #dict where key is class, value is list of ground truth boxes for that class\n",
    "        gt_boxes_by_class = {k:[g for g in gt_boxes if g[5] == k] for k in range(self.kwargs[\"num_classes\"]) }\n",
    "        \n",
    "        #dict where key is class, value is list of confidences for each box guess\n",
    "        pred_confs = {k:[] for k in range(self.kwargs[\"num_classes\"]) }\n",
    "        \n",
    "        #dict where key is class, value is list of binary variables, where 1 means, there is a ground truth box\n",
    "        #for the corresponding guess in pred_confs\n",
    "        gt_confs = {k:[] for k in range(self.kwargs[\"num_classes\"])}\n",
    "        \n",
    "\n",
    " \n",
    "        gt_boxes_byc = deepcopy(gt_boxes_by_class)\n",
    "\n",
    "        #for each predicted box (in order from highest conf to lowest)\n",
    "        for c in pred_boxes_by_class.keys():\n",
    "            for pc_box in pred_boxes_by_class[c]:\n",
    "                conf = pc_box[4]\n",
    "                cls = pc_box[5]\n",
    "                ious = np.asarray([self.iou(pc_box, gc_box) for gc_box in gt_boxes_byc[c] ])\n",
    "                \n",
    "                # get all gt boxes that have iou with given box over the threshold\n",
    "                C = [ind for ind, gc_box in enumerate(gt_boxes_byc[c]) if ious[ind] > self.kwargs[\"iou_thresh\"] ]\n",
    "                \n",
    "                #this means there is at least once gt box that overlaps over \n",
    "                #the IOU threshold with predicted\n",
    "                if len(C) > 0:\n",
    "                    # if there are some gt ones that are over the threshold\n",
    "                    # grab the highest one in terms of iou\n",
    "                    max_ind = np.argmax(ious)\n",
    "                    \n",
    "                    \n",
    "                    # remove this box from consideration\n",
    "                    del gt_boxes_byc[c][max_ind]\n",
    "                    \n",
    "                    #add the confidence to the prediction\n",
    "                    pred_confs[c].append(conf)\n",
    "                    #add a 1 to the ground truth\n",
    "                    gt_confs[c].append(1)\n",
    "                \n",
    "                #no box overlaps with the guess. Put a zero for ground truth here\n",
    "                else:\n",
    "                    pred_confs[cls].append(conf)\n",
    "                    gt_confs[c].append(0)\n",
    "        \n",
    "\n",
    "        return pred_confs, gt_confs\n",
    "\n",
    "\n",
    "\n",
    "    \n",
    "    def get_all_boxes(self, pred_tensor,y_tensor):\n",
    "        all_gt_boxes = []\n",
    "        all_pred_boxes = []\n",
    "        for i in range(pred_tensor.shape[0]):\n",
    "            pred_boxes, gt_boxes = self.get_boxes(pred_tensor[i], y_tensor[i])\n",
    "            all_gt_boxes.append(gt_boxes)\n",
    "            all_pred_boxes.append(pred_boxes)\n",
    "        return all_pred_boxes, all_gt_boxes\n",
    "\n",
    "\n",
    "    def get_boxes(self, pred_tensor, y_tensor):\n",
    "        ''' pred_tensor is of shape: 5 + num_classes, x_g, y_g \n",
    "            y_tensor is same shape'''\n",
    "        '''returns two dicts of the form\n",
    "             key (class), value: list of boxes with conf for that class'''\n",
    "        pred_boxes = self.get_boxes_from_tensor(pred_tensor)\n",
    "        gt_boxes = self.get_boxes_from_tensor(y_tensor)\n",
    "        #print len(gt_boxes)\n",
    "        return pred_boxes, gt_boxes\n",
    "\n",
    "    def get_boxes_from_tensor(self, tensor):\n",
    "        #tensor is numpy tensor\n",
    "        x_g, y_g = tensor.shape[-2], tensor.shape[-1]\n",
    "        boxes = []\n",
    "\n",
    "        # pull out all box guesses\n",
    "        for i in range(x_g):\n",
    "            for j in range(y_g):\n",
    "                #print tensor.shape\n",
    "                coords = tensor[:5,i,j]\n",
    "                box = self.convert_coords_to_box(coords,i,j)\n",
    "                conf = coords[-1]\n",
    "                cls = np.argmax(tensor[6:,i,j]) # + 1 # cuz classes go from 1 to4\n",
    "                box.append(cls)\n",
    "                #if conf > self.kwargs[\"conf_thresh\"]:\n",
    "                boxes.append(box)\n",
    "\n",
    "        #sort by confidence\n",
    "        boxes.sort(lambda a,b: -1 if a[4] > b[4] else 1)        \n",
    "        return boxes\n",
    "\n",
    "\n",
    "    def convert_coords_to_box(self, coords, xind, yind):\n",
    "        scale_factor = self.kwargs[\"scale_factor\"]\n",
    "        #print coords\n",
    "        xoff,yoff,w,h, conf = coords\n",
    "        x,y = scale_factor*(xind + xoff), scale_factor *(yind + yoff)\n",
    "\n",
    "        w,h = 2**w * scale_factor, 2**h * scale_factor\n",
    "\n",
    "        return [x,y,w,h,conf]\n",
    "\n",
    "    def iou(self, box1,box2):\n",
    "        #box1 and box2 are numpy arrays\n",
    "        #boxes are expected in x_center, y_center, width, height format\n",
    "        x1,y1,w1,h1 = box1[:4]\n",
    "        x2,y2,w2,h2 = box2[:4]\n",
    "        #print w1,h1\n",
    "        #print w2,h2\n",
    "        xmin1, xmax1, ymin1, ymax1 = max(0, x1 - w1 / 2.), x1 + w1 /2., max(0,y1 - h1 / 2.), y1 + h1 /2.\n",
    "        xmin2, xmax2, ymin2, ymax2 = max(0,x2 - w2 / 2.), x2 + w2 /2,max(0,y2 - h2 / 2), y2 + h2 /2\n",
    "        inters = max(0,(min(xmax1,xmax2) - max(xmin1,xmin2)))   * \\\n",
    "                              max(0,(min(ymax1,ymax2) - max(ymin1,ymin2)) )\n",
    "        def get_area(box_mm):\n",
    "            xmin, xmax, ymin, ymax = box_mm\n",
    "            area = (xmax - xmin) * (ymax - ymin)\n",
    "    #         if area == 0.0:\n",
    "    #             print \"aaaaaah\", xmin, xmax, ymin, ymax\n",
    "            return area\n",
    "        union = get_area((xmin1, xmax1, ymin1, ymax1)) + get_area((xmin2, xmax2, ymin2, ymax2)) - inters                                                        \n",
    "\n",
    "        return inters / float(union)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_detec_loss(pred, gt, kwargs):\n",
    "    #TODO add in multiple bbox behavior\n",
    "    #pred is n_ex, [x,y,w,h,c,classes], x, y\n",
    "    #get number of examples and the indices of the tesnor \n",
    "    #to where x,y coirrds height width and confidence go\n",
    "    pred = pred.transpose((0,2,3,1))\n",
    "    gt = gt.transpose((0,2,3,1))\n",
    "    nex = pred.shape[0]\n",
    "    cinds = T.arange(6)\n",
    "    \n",
    "    #x coord indices, y coord indices, width, height, confidence\n",
    "    xs,ys,ws,hs,cs, csn = cinds[0], cinds[1], cinds[2], cinds[3], cinds[4], cinds[5]\n",
    "    \n",
    "    #index for prob vector for all classes\n",
    "    ps = T.arange(6,pred.shape[3])\n",
    "    \n",
    "    #theano will now make elements less than or equal to 0 as zero and others 1 (so output shape is )\n",
    "    obj_inds = gt[:,:, :,cs] > 0.\n",
    "   \n",
    "    #use nonzero in order to get boolean indexing  (eliminate the indices that are zero)\n",
    "    #get specific x,y location of gt objects and the predicted output for that x,y location\n",
    "    tg_obj = gt[obj_inds.nonzero()]\n",
    "    tp_obj = pred[obj_inds.nonzero()]\n",
    "    \n",
    "    #term1\n",
    "    #take the sum of squared difference between predicted and gt for the x and y corrdinate \n",
    "    s_x = smoothL1(tp_obj[:,xs] - tg_obj[:,xs])\n",
    "\n",
    "    s_y = smoothL1(tp_obj[:,ys] - tg_obj[:,ys])\n",
    "\n",
    "    raw_loss1 = T.sum(s_x + s_y)\n",
    "\n",
    "    #multipily by lambda coord (the scaling factor for bbox coords)\n",
    "    sterm1 = kwargs['coord_penalty'] * raw_loss1\n",
    "\n",
    "\n",
    "    #term2\n",
    "\n",
    "    #get sum of squared diff of the of heights and widths b/w pred and gt normalized by squared heights and widths of gt \n",
    "    s_w = smoothL1(tp_obj[:,ws] - tg_obj[:,ws]) #/ T.square(tg_obj[:,ws])\n",
    "    s_h = smoothL1((tp_obj[:,hs] - tg_obj[:,hs])) #/ T.square(tg_obj[:,hs])\n",
    "    raw_loss2 = T.sum(s_w + s_h)\n",
    "\n",
    "    sterm2 = kwargs['size_penalty'] * raw_loss2\n",
    "\n",
    "\n",
    "    #term3\n",
    "    #get sum of squared diff between confidence for places with actual boxes of pred vs. ground truth\n",
    "    s_c  = -T.log(tp_obj[:,cs]+ 0.00001)\n",
    "    raw_loss3 = T.sum(s_c)\n",
    "    sterm3 = raw_loss3\n",
    "\n",
    "\n",
    "    #term4\n",
    "    #get the real coordinates where there are no objects\n",
    "    no_ind  = gt[:,:,:,cs] <= 0.\n",
    "    tp_no_obj = pred[no_ind.nonzero()]\n",
    "\n",
    "    #get the log likelhood that there isn't a box\n",
    "    s_nc = -T.log(tp_no_obj[:,csn]+ 0.00001)\n",
    "\n",
    "    raw_loss4 = T.sum(s_nc)\n",
    "\n",
    "    sterm4 = kwargs['nonobj_penalty'] * raw_loss4\n",
    "\n",
    "\n",
    "    #get the cross entropy of these softmax vectors\n",
    "    s_p = T.nnet.categorical_crossentropy(tp_obj[:,ps], tg_obj[:,ps])\n",
    "\n",
    "    raw_loss5 = T.sum(s_p)\n",
    "    sterm5 = raw_loss5\n",
    "\n",
    "    #adds up terms divides by number of examples in the batch\n",
    "    loss = (1. / nex) * (sterm1 + sterm2 + sterm3 + sterm4 + sterm5)\n",
    "    return loss, [sterm1, sterm2, sterm3, sterm4, sterm5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def nms(boxes):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def dump_hyperparams(dic, path):\n",
    "    new_dic = {k:str(dic[k]) for k in dic.keys()}\n",
    "    with open(path + '/hyperparams.txt', 'w') as f:\n",
    "        for k,v in new_dic.iteritems():\n",
    "            f.write(k + ' : ' + v + \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "def setup_logging(save_path):\n",
    "    logger = logging.getLogger('simple_example')\n",
    "    logger.setLevel(logging.DEBUG)\n",
    "    # create file handler which logs even debug messages\n",
    "    fh = logging.FileHandler('%s/training.log'%(save_path))\n",
    "    fh.setLevel(logging.DEBUG)\n",
    "    # create console handler with a higher log level\n",
    "    ch = logging.StreamHandler()\n",
    "    ch.setLevel(logging.DEBUG)\n",
    "    logger.addHandler(ch)\n",
    "    logger.addHandler(fh)\n",
    "    return logger"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [conda root]",
   "language": "python",
   "name": "conda-root-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
